#include <algorithm>
#include <cuda_runtime.h>
#include <mma.h>
#include <vector>
#include <iostream>
#include <cmath>
#include <cuda_fp16.h>
#include <iomanip> 
#include <chrono>
using namespace nvcuda;

const int WMMA_M = 16;
const int WMMA_N = 16;
const int WMMA_K = 16;

// 矩阵填充函数
void padMatrix(std::vector<std::vector<float>>& mat, int target_rows, int target_cols) {
    for (auto& row : mat) {
        row.resize(target_cols, 0.0f);
    }
    while (mat.size() < target_rows) {
        mat.push_back(std::vector<float>(target_cols, 0.0f));
    }
}

__global__ void matrixMultiplyWMMA(__half* A, __half* B, float* C, int M, int N, int K) {
    int col = blockIdx.x * WMMA_M + threadIdx.x;
    int row = blockIdx.y * WMMA_N + threadIdx.y;

    if (row < M && col < N) {
        wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, half, wmma::col_major> a_frag;
        wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, half, wmma::row_major> b_frag;
        wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float> c_frag;

        wmma::fill_fragment(c_frag, 0.0f);

        for (int i = 0; i < K; i += WMMA_K) {
            wmma::load_matrix_sync(a_frag, A + row * K + i, K);
            wmma::load_matrix_sync(b_frag, B + i * N + col, N);
            wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
        }

        wmma::store_matrix_sync(C + row * N + col, c_frag, N, wmma::mem_row_major);
    }
}

void convertFp32ToFp16(std::vector<std::vector<float>>& input, __half* output, int rows, int cols) {
    for (int i = 0; i < rows; ++i) {
        for (int j = 0; j < cols; ++j) {
            output[i * cols + j] = __float2half(input[i][j]);
        }
    }
}

void matrixMultiply(const std::vector<std::vector<float>>& A, const std::vector<std::vector<float>>& B, std::vector<std::vector<float>>& C, int M, int N, int K) {
    int padded_M = ((M + WMMA_M - 1) / WMMA_M) * WMMA_M;
    int padded_N = ((N + WMMA_N - 1) / WMMA_N) * WMMA_N;
    int padded_K = ((K + WMMA_K - 1) / WMMA_K) * WMMA_K;

    std::vector<std::vector<float>> padded_A = A;
    std::vector<std::vector<float>> padded_B = B;
    
    padMatrix(padded_A, padded_M, padded_K);
    padMatrix(padded_B, padded_K, padded_N);
    C.resize(padded_M, std::vector<float>(padded_N, 0.0f));

    __half *d_A, *d_B;
    float *d_C;
    cudaMalloc(&d_A, padded_M * padded_K * sizeof(__half));
    cudaMalloc(&d_B, padded_K * padded_N * sizeof(__half));
    cudaMalloc(&d_C, padded_M * padded_N * sizeof(float));

    __half* h_A = new __half[padded_M * padded_K];
    __half* h_B = new __half[padded_K * padded_N];

    convertFp32ToFp16(padded_A, h_A, padded_M, padded_K);
    convertFp32ToFp16(padded_B, h_B, padded_K, padded_N);
    
    cudaMemcpy(d_A, h_A, padded_M * padded_K * sizeof(__half), cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, padded_K * padded_N * sizeof(__half), cudaMemcpyHostToDevice);
    dim3 threadsPerBlock(WMMA_M, WMMA_N);
    dim3 numBlocks(padded_N / WMMA_N, padded_M / WMMA_M);
    matrixMultiplyWMMA<<<numBlocks, threadsPerBlock>>>(d_A, d_B, d_C, padded_M, padded_N, padded_K);
    cudaError_t error = cudaGetLastError();
if (error != cudaSuccess) {
    std::cerr << "CUDA error: " << cudaGetErrorString(error) << std::endl;
}
cudaDeviceSynchronize(); 
    std::vector<float> temp(padded_M * padded_N);  // 创建临时的连续内存数组
cudaMemcpy(temp.data(), d_C, padded_M * padded_N * sizeof(float), cudaMemcpyDeviceToHost);  // 从设备到临时数组


// 将数据从临时数组复制回C的二维结构
for (int i = 0; i < M; ++i) {
    std::copy_n(temp.begin() + i * N, N, C[i].begin());
}
    // Remove padding from the result
    C.resize(M);
    for (auto& row : C) {
        row.resize(N);
    }

    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C);
    delete[] h_A;
    delete[] h_B;
}



int main(){
    const int size = 64;  // Set the size of the matrix to 16x16
    int M,N,K;
    M=N=K=size;
    std::vector<std::vector<float>> A(size, std::vector<float>(size));
    std::vector<std::vector<float>> B(size, std::vector<float>(size));
    std::vector<std::vector<float>> C(size, std::vector<float>(size));
    // Fill matrices with some values
    for (int i = 0; i < size; ++i) {
        for (int j = 0; j < size; ++j) {
            // A[i][j] = 1;  // Simple pattern for easy verification
            // if (i>j)
            //     B[i][j]=1;
            // else if (j>i)
            //     B[i][j]=-1;
            // else 
            //     B[i][j]=0;
            A[i][j] = 1;//i%100 *0.01;
            B[i][j] = 1;//j%100 *0.01;
        }
    }
    auto startWMMA = std::chrono::high_resolution_clock::now();
    matrixMultiply(A, B, C, M, N, K);
    cudaDeviceSynchronize(); // 等待GPU完成操作
    auto endWMMA = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double, std::milli> wmmaTime = endWMMA - startWMMA;
    std::cout << "WMMA Time: " << wmmaTime.count() << " ms" << std::endl;
    std::cout<< std::endl;
}